# !/usr/bin/env python
# -*- coding: utf-8 -*-

"""
__project__ =  'tricker'
__file__    =  'tasks.py'
__author__  =  'king'
__time__    =  '2024/5/29 上午10:20'


                              _ooOoo_
                             o8888888o
                             88" . "88
                             (| -_- |)
                             O\  =  /O
                          ____/`---'\____
                        .'  \\|     |//  `.
                       /  \\|||  :  |||//  \
                      /  _||||| -:- |||||-  \
                      |   | \\\  -  /// |   |
                      | \_|  ''\---/''  |   |
                      \  .-\__  `-`  ___/-. /
                    ___`. .'  /--.--\  `. . __
                 ."" '<  `.___\_<|>_/___.'  >'"".
                | | :  `- \`.;`\ _ /`;.`/ - ` : | |
                \  \ `-.   \_ __\ /__ _/   .-` /  /
           ======`-.____`-.___\_____/___.-`____.-'======
                              `=---='
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
                       佛祖保佑        永无BUG
"""
import os
import shutil
import jinja2
import psutil
import tempfile
from pathlib import Path
from taskflow import task, engines
from taskflow.patterns import linear_flow

from tricker.commands.image.constants import (
    INIT_PACKAGES,
    INSTALLED_PACKAGES,
    REMOVED_PACKAGES,
)
from tricker.common import utils
from tricker.common.log import LOG
from tricker.commands.image.exceptions import MemoryNotEnoughException


def get_dev_pts_options() -> str:
    for partitions in psutil.disk_partitions(all=True):
        if partitions.mountpoint == '/dev/pts':
            return partitions.opts


def get_templates_path() -> str:
    path = os.path.abspath(__file__)
    path = os.path.dirname(path)
    return f'{path}/templates'


class CheckLocalMemoryTask(task.Task):
    def execute(self, *args, **kwargs):
        memory = psutil.virtual_memory()
        require_memory = 4 * 2 * 1024 * 1024 * 1024
        if memory.total < require_memory:
            LOG.error(
                'Not enough RAM to use tmpfs for build, require 8 GB at least'
            )
            raise MemoryNotEnoughException()


class CreateRepoTask(task.Task):
    def __init__(self):
        super().__init__()

    def execute(self, *args, **kwargs):
        utils.generate_repo()


class ImageWorker(object):
    def __init__(self):
        self.builder_path, self.image_path = self.create_build_path()
        self.repo_path = f'{self.builder_path}/yum.repos.d/'
        self.mount_path = f'{self.builder_path}/mnt/'
        self.proc_path = f'{self.mount_path}/proc/'
        self.dev_path = f'{self.mount_path}/dev/'
        self.sys_path = f'{self.mount_path}/sys/'
        self.rpm_path = f'{self.mount_path}/var/lib/rpm/'
        self.etc_path = f'{self.mount_path}/etc/'
        self.tmp_yum_path = f'{self.mount_path}/tmp/yum/'
        home_path = os.path.expanduser('~')
        self.cache_path = f'{home_path}/.cache/create_image/'
        self.yum_cache_path = f'{self.cache_path}/yum/'

    @staticmethod
    def create_build_path() -> (str, str):
        date = utils.formatted_date()
        builder_dir = f'/tmp/tricker_builder.{date}'
        image_dir = f'/tmp/tricker_image.{date}'
        Path(builder_dir).mkdir(parents=True, exist_ok=True)
        Path(image_dir).mkdir(parents=True, exist_ok=True)
        uid, gid = os.getuid(), os.getgid()
        os.chown(builder_dir, uid, gid)
        os.chown(image_dir, uid, gid)
        return builder_dir, image_dir

    def create_repo_file(self, path: str):
        with open(f'{get_templates_path()}/openeuler.repo.j2', 'r+') as f:
            data = f.read()
            template = jinja2.Template(data).render()
            with open(f'{path}/openeuler.repo', 'w+') as f:
                f.writelines(template)

    @staticmethod
    def _run():
        flow = linear_flow.Flow('image_builder')
        flow.add(CheckLocalMemoryTask())
        e = engines.load(flow, engine='parallel', max_workers=1)
        e.run()

    @staticmethod
    def check_memory():
        memory = psutil.virtual_memory()
        require_memory = 4 * 2 * 1024 * 1024 * 1024
        if memory.total < require_memory:
            LOG.error(
                'Not enough RAM to use tmpfs for build, require 8 GB at least'
            )
            raise MemoryNotEnoughException()

    def create_path(self):
        paths = [
            self.repo_path,
            self.proc_path,
            self.dev_path,
            self.sys_path,
            self.rpm_path,
            f'{self.etc_path}/yum.repos.d/',
            self.tmp_yum_path,
            self.yum_cache_path,
        ]
        for path in paths:
            Path(path).mkdir(parents=True, exist_ok=True)

    def mount_devices(self):
        options = get_dev_pts_options()
        commands = [
            f'sudo mount -t proc none {self.proc_path}',
            f'sudo mount --bind /dev {self.dev_path}',
            f'sudo mount -t devpts -o {options} devpts {self.dev_path}/pts',
            f'sudo mount -o ro -t sysfs none {self.sys_path}',
            f'sudo mount --bind {self.yum_cache_path} {self.tmp_yum_path}',
        ]
        for command in commands:
            utils.execute(command)

    def init_rpm(self):
        utils.execute_rpm(self.mount_path, '--initdb')

    def copy_files(self):
        files = [
            '/etc/resolv.conf',
        ]
        for f in files:
            shutil.copy(f, f'{self.mount_path}/{f}')

    def install_repo(self):
        packages = ' '.join(
            [
                'basesystem',
                'filesystem',
                'setup',
                'openEuler-release',
                'openEuler-repos',
                'openEuler-gpg-keys',
            ]
        )
        cache_dir = tempfile.TemporaryDirectory()
        save_dir = tempfile.TemporaryDirectory()
        command = (
            f'dnf download --verbose '
            f'--setopt=reposdir={self.repo_path} '
            f'--setopt=cachedir={cache_dir.name} '
            f'--destdir={save_dir.name} {packages}'
        )
        utils.execute(command)
        utils.execute_rpm(
            self.mount_path, f'--nodeps -ivh {save_dir.name}/*rpm'
        )
        cache_dir.cleanup()
        save_dir.cleanup()

    def install_packages(self):
        packages = ' '.join(INIT_PACKAGES)
        command = (
            f'sudo -E dnf -y '
            f'--disableexcludes=all '
            f'--setopt=cachedir={self.yum_cache_path} '
            f'--setopt=reposdir={self.mount_path}/etc/yum.repos.d '
            f'--installroot {self.mount_path} '
            f'install {packages}'
        )
        utils.execute(command)

    def run_chroot(self):
        packages = ' '.join(INSTALLED_PACKAGES)
        remove = ''.join(REMOVED_PACKAGES)
        commands = [
            'rpm --rebuilddb',
            'dnf clean all',
            'dnf -y update',
            f'dnf install -y '
            f'--setopt=cachedir={self.yum_cache_path} {packages}',
            'ln -sf /usr/share/zoneinfo/UTC /etc/localtime',
        ]
        for command in commands:
            LOG.info('executing %s in chroot' % command)
            utils.execute_chroot(self.mount_path, command)

    def write_files(self):
        file_content_mapping = {
            '/etc/resolv.conf': 'nameserver 8.8.8.8',
            '/etc/locale.conf': 'LANG="en_US.UTF-8"',
        }
        for file_path, content in file_content_mapping.items():
            LOG.info('writing file %s' % file_path)
            with open(f'{self.mount_path}/{file_path}', 'w') as f:
                f.write(content)

    def umount_devices(self):
        paths = [
            self.tmp_yum_path,
            self.proc_path,
            f'{self.dev_path}/pts',
            self.dev_path,
            self.sys_path,
        ]
        for path in paths:
            LOG.info('umount path %s' % path)
            LOG.info(f'sudo umount {path}')
            # utils.execute(f'sudo umount {path}')

    def run(self):
        self.check_memory()
        self.create_path()
        self.create_repo_file(self.repo_path)
        self.mount_devices()
        self.init_rpm()
        self.copy_files()
        self.install_repo()
        self.install_packages()
        self.run_chroot()
        self.write_files()
        self.umount_devices()
